package com.gitee.helioz.openai.demo;

import com.gitee.helioz.openai.api.OpenAiAHCServiceImpl;
import com.gitee.helioz.openai.api.OpenAiHttpClientServiceImpl;
import com.gitee.helioz.openai.api.OpenAiOkHttpServiceImpl;
import com.gitee.helioz.openai.api.OpenAiService;
import com.gitee.helioz.openai.api.OpenAiURLConnectionServiceImpl;
import com.gitee.helioz.openai.api.dto.DeleteResponse;
import com.gitee.helioz.openai.api.dto.audio.TranscriptionsRequest;
import com.gitee.helioz.openai.api.dto.audio.TranscriptionsResponse;
import com.gitee.helioz.openai.api.dto.audio.TranslationsRequest;
import com.gitee.helioz.openai.api.dto.chat.ChatCompletionRequest;
import com.gitee.helioz.openai.api.dto.chat.ChatCompletionResponse;
import com.gitee.helioz.openai.api.dto.embeddings.EmbeddingRequest;
import com.gitee.helioz.openai.api.dto.embeddings.EmbeddingResponse;
import com.gitee.helioz.openai.api.dto.files.FileDTO;
import com.gitee.helioz.openai.api.dto.files.FileResponse;
import com.gitee.helioz.openai.api.dto.files.FileUploadRequest;
import com.gitee.helioz.openai.api.dto.finetune.FineTuneCreateRequest;
import com.gitee.helioz.openai.api.dto.finetune.FineTuneEventsResponse;
import com.gitee.helioz.openai.api.dto.finetune.FineTuneListResponse;
import com.gitee.helioz.openai.api.dto.finetune.FineTuneResponse;
import com.gitee.helioz.openai.api.dto.images.ImageEditRequest;
import com.gitee.helioz.openai.api.dto.images.ImageGenerationRequest;
import com.gitee.helioz.openai.api.dto.images.ImageGenerationResponse;
import com.gitee.helioz.openai.api.dto.images.ImageVariationRequest;
import com.gitee.helioz.openai.api.dto.images.SizeEnum;
import com.gitee.helioz.openai.api.dto.models.ModelDTO;
import com.gitee.helioz.openai.api.dto.models.ModelListResponse;
import com.gitee.helioz.openai.api.dto.moderations.ModerationCreateRequest;
import com.gitee.helioz.openai.api.dto.moderations.ModerationCreateResponse;
import com.gitee.helioz.openai.common.authentication.AuthenticationConfigDTO;
import com.gitee.helioz.openai.common.constants.MessageRoleEnum;
import com.gitee.helioz.openai.common.util.http.HttpFilePart;
import com.gitee.helioz.openai.common.util.json.GsonUtil;
import com.google.common.collect.Lists;
import com.google.common.io.Files;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;

import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;

/**
 * 〈〉
 *
 * @author <a href="https://gitee.com/helioz">Morning Starr</a>
 */
@Slf4j
public class OpenAiServiceDemo extends BaseAppTest {

//    private static final OpenAiService openAiService = new OpenAiURLConnectionServiceImpl();
//    private static final OpenAiService openAiService = new OpenAiAHCServiceImpl();
//    private static final OpenAiService openAiService = new OpenAiHttpClientServiceImpl();
    private static final OpenAiService openAiService = new OpenAiOkHttpServiceImpl();

    @BeforeAll
    public static void initService() throws IOException {
        log.info("装载 API KEY");
        String filePath = "F:\\学习\\openai\\api_keys.txt";
        List<String> apiKeyLines = Files.readLines(new File(filePath), StandardCharsets.UTF_8);
        for (String apiKeyLine : apiKeyLines) {
            String[] apiArr = StringUtils.split(apiKeyLine, "||");
            AuthenticationConfigDTO authConfig = AuthenticationConfigDTO.builder()
                    .apiKeyName(apiArr[0])
                    .apiKey(apiArr[1])
                    .build();
            log.debug("API KEY：" + authConfig);
            openAiService.addAuthConfig(authConfig.getApiKey(), authConfig);
        }
    }

    // region models test
    @Test
    @DisplayName("List models")
    void testModelsList() {
        ModelListResponse response = openAiService.getModelsService().list();
        log.info(response.getObject());
        log.info("列出当前可用的模型结果={}", response);
    }

    @Test
    @DisplayName("Retrieve model")
    void testModelsRetrieve() throws InterruptedException {
        CountDownLatch countDownLatch = new CountDownLatch(1);
        new Thread(() -> {
            openAiService.getAuthConfig();
            countDownLatch.countDown();
        }).start();
        countDownLatch.await();
        ModelDTO response = openAiService.getModelsService().retrieve("text-davinci-003");
        log.info(response.getObject());
        log.info("检索模型实例结果={}", response);
    }
    // endregion

    // region completions test
    @Test
    @DisplayName("Create completion")
    void testCompletionCreate() {
//        ModelDTO response = openAiService.getCompletionService().create("text-davinci-003");
//        log.info(response.getObject());
//        log.info("检索模型实例结果={}", response);
    }
    // endregion

    // region chat test
    @Test
    @DisplayName("Create chat completion")
    void testChatCreate() {
        ChatCompletionRequest ccRequest = new ChatCompletionRequest();
        ccRequest.setModel("gpt-3.5-turbo");
        ccRequest.setUser("Morning Starr");
        List<ChatCompletionRequest.ChatMessage> messages = new ArrayList<>();
        messages.add(ChatCompletionRequest.ChatMessage.builder().role(MessageRoleEnum.USER.getType()).content("请问天空是蓝色还是黑色？").build());
        ccRequest.setMessages(messages);
        ChatCompletionResponse response = openAiService.getChatService().createChatCompletion(ccRequest);
        log.info(response.getObject());
        log.info("创建聊天式对话补全请结果={}", GsonUtil.toJson(response));
    }
    // endregion

    // region edits test
    @Test
    @DisplayName("Create edit")
    void testEditsCreate() {
//        ModelDTO response = openAiService.getEditService().create("text-davinci-003");
//        log.info(response.getObject());
//        log.info("检索模型实例结果={}", response);
    }
    // endregion

    // region images test
    @Test
    @DisplayName("Create image")
    void testImagesCreate() {
        ImageGenerationRequest generationRequest = new ImageGenerationRequest();
        generationRequest.setPrompt("红色的皮卡丘");
        ImageGenerationResponse response = openAiService.getImageService().generations(generationRequest);
        log.info("生成图像结果={}", response);
    }

    @Test
    @DisplayName("Create image edit")
    void testImagesEdit() throws IOException {
        ImageEditRequest editRequest = new ImageEditRequest();
        HttpFilePart image = new HttpFilePart();
        File imageFile = new File("");
        image.setFileName(imageFile.getName());
        image.setFileBytes(Files.toByteArray(imageFile));
        editRequest.setImage(image);
        editRequest.setPrompt("红色的皮卡丘，释放着十万伏特的电");
        ImageGenerationResponse response = openAiService.getImageService().edits(editRequest);
        log.info("编辑图像结果={}", response);
    }

    @Test
    @DisplayName("Create image variation")
    void testImagesVariation() throws IOException {
        ImageVariationRequest variationRequest = new ImageVariationRequest();
        HttpFilePart image = new HttpFilePart();
        File imageFile = new File("");
        image.setFileName(imageFile.getName());
        image.setFileBytes(Files.toByteArray(imageFile));
        variationRequest.setImage(image);
        variationRequest.setSize(SizeEnum.size_512.getType());
        ImageGenerationResponse response = openAiService.getImageService().variations(variationRequest);
        log.info("创建图像变体结果={}", response);
    }
    // endregion

    // region embeddings test
    @Test
    @DisplayName("Create embeddings")
    void testEmbeddingsCreate() {
        EmbeddingRequest embeddingRequest = new EmbeddingRequest();
        embeddingRequest.setModel("text-similarity-babbage-001");
        embeddingRequest.setInput(Lists.newArrayList("教室教室土木三三二二", "教师教师干木三三二二"));
        EmbeddingResponse response = openAiService.getEmbeddingService().create(embeddingRequest);
        log.info("获取给定输入的向量表示结果={}", response);
    }
    // endregion

    // region audio test
    @Test
    @DisplayName("Create transcription")
    void testAudioTranscriptions() throws IOException {
        TranscriptionsRequest transRequest = new TranscriptionsRequest();
        HttpFilePart audio = new HttpFilePart();
        File audioFile = new File("");
        audio.setFileName(audioFile.getName());
        audio.setFileBytes(Files.toByteArray(audioFile));
        transRequest.setFile(audio);
        TranscriptionsResponse response = openAiService.getAudioService().transcriptions(transRequest);
        log.info("将音频文件转录为文本结果={}", response);
    }

    @Test
    @DisplayName("Create translation")
    void testAudioTranslations() throws IOException {
        TranslationsRequest transRequest = new TranslationsRequest();
        HttpFilePart audio = new HttpFilePart();
        File audioFile = new File("");
        audio.setFileName(audioFile.getName());
        audio.setFileBytes(Files.toByteArray(audioFile));
        transRequest.setFile(audio);
        TranscriptionsResponse response = openAiService.getAudioService().translations(transRequest);
        log.info("将音频翻译成英语结果={}", response);
    }
    // endregion

    // region files test
    @Test
    @DisplayName("List files")
    void testFilesList() {
        FileResponse response = openAiService.getFileService().list();
        log.info("获取属于其组织的文件列表结果={}", response);
    }

    @Test
    @DisplayName("Upload file")
    void testFilesUpload() throws IOException {
        FileUploadRequest fileRequest = new FileUploadRequest();
        HttpFilePart file = new HttpFilePart();
        File textFile = new File("F:\\学习\\openai\\test_file_upload.jsonl");
        file.setFileName(textFile.getName());
        file.setFileBytes(Files.toByteArray(textFile));
        fileRequest.setFile(file);
        fileRequest.setPurpose("fine-tune");
        FileDTO response = openAiService.getFileService().upload(fileRequest);
        log.info("上传文件结果={}", response);
    }

    @Test
    @DisplayName("Delete file")
    void testFilesDelete() {
        DeleteResponse response = openAiService.getFileService().delete("file-lmPCxr3ON98XYqzp7sEB5VWV");
        log.info("删除文件结果={}", response);
    }

    @Test
    @DisplayName("Retrieve file")
    void testFilesRetrieve() {
        FileDTO response = openAiService.getFileService().retrieve("file-lmPCxr3ON98XYqzp7sEB5VWV");
        log.info("获取文件信息结果={}", response);
    }

    @Test
    @DisplayName("Retrieve file content\n")
    void testFilesRetrieveContent() {
        String response = openAiService.getFileService().retrieveContent("file-8iAzZpBWS2OHOtLG6Dw0G3bb");
        log.info("获取文件内容结果={}", response);
    }
    // endregion

    // region fine-tunes test
    @Test
    @DisplayName("Create fine-tune")
    void testFineTunesCreate() {
        FineTuneCreateRequest createRequest = new FineTuneCreateRequest();
        createRequest.setTrainingFile("file-8iAzZpBWS2OHOtLG6Dw0G3bb");
        FineTuneResponse response = openAiService.getFineTuneService().create(createRequest);
        // ft-BbhRcKcxe7aVE0MSINgxiFWK
        log.info("创建一个用于微调（fine-tune）模型的训练任务结果={}", response);
    }

    @Test
    @DisplayName("List fine-tunes\n")
    void testFineTunesList() {
        FineTuneListResponse response = openAiService.getFineTuneService().list();
        log.info("微调（fine-tune）模型的训练任务集合结果={}", response);
    }

    @Test
    @DisplayName("Retrieve fine-tune")
    void testFineTunesRetrieve() {
        FineTuneResponse response = openAiService.getFineTuneService().retrieve("ft-BbhRcKcxe7aVE0MSINgxiFWK");
        log.info("获取微调（fine-tune）模型的训练任务信息结果={}", response);
    }

    @Test
    @DisplayName("Cancel fine-tune")
    void testFineTunesCancel() {
        FineTuneResponse response = openAiService.getFineTuneService().cancel("ft-BbhRcKcxe7aVE0MSINgxiFWK");
        log.info("取消微调（fine-tune）模型的训练任务结果={}", response);
    }

    @Test
    @DisplayName("List fine-tune events")
    void testFineTunesEvents() {
        // TODO: 流式
        FineTuneEventsResponse response = openAiService.getFineTuneService().events("ft-BbhRcKcxe7aVE0MSINgxiFWK", true);
        log.info("获取指定微调任务的细粒度状态更新结果={}", response);
    }

    @Test
    @DisplayName("Delete fine-tune model")
    void testFineTunesDelete() {
        DeleteResponse response = openAiService.getFineTuneService().delete("curie");
        log.info("删除微调（fine-tune）模型结果={}", response);
    }
    // endregion

    // region moderations test
    @Test
    @DisplayName("Create moderation")
    void testModerationsCreate() {
        ModerationCreateRequest createRequest = new ModerationCreateRequest();
        createRequest.setInput(Lists.newArrayList());
        createRequest.setModel("text-moderation-stable");
        ModerationCreateResponse response = openAiService.getModerationService().create(createRequest);
        // ft-BbhRcKcxe7aVE0MSINgxiFWK
        log.info("创建内容过滤（Moderation）任务结果={}", response);
    }
    // endregion

}
